package jp.aist.gtrc.plus.admission.EDF;

import java.text.DateFormat;
import java.util.*;

public class Schedule {
   // Expected start time assumes the execution times of all jobs 
   //  given by users are correct.
   private Map<Job,Integer> expectedStartTimeMap; // map job -> time 
   private Map<String,List<Job>> jobQueueMap; // map CPU -> queue
   private Map<String,Set<String>> nodeQueueMap; // map node -> set of CPUs

   private Map<String,Integer> savedJobQueueSizeMap;   // map CPU -> size
   private Map<Job,Integer> savedExpectedStartTimeMap; 

   public Schedule() {
      expectedStartTimeMap = new HashMap<Job,Integer>();
      jobQueueMap = new HashMap<String,List<Job>>();
      nodeQueueMap = new HashMap<String,Set<String>>();
   }

   public synchronized void appendQueue(CPU cpu, Job job, int st) {
      Set<String> cpus = nodeQueueMap.get(cpu.node);
      if (cpus == null) {
         cpus = new HashSet<String>();
         nodeQueueMap.put(cpu.node,cpus);
      }
      cpus.add(cpu.id);
      List<Job> jobQueue = jobQueueMap.get(cpu.id);
      if (jobQueue == null) {
         jobQueue = new ArrayList<Job>();
         jobQueueMap.put(cpu.id,jobQueue);
      }
      jobQueue.add(job);
      expectedStartTimeMap.put(job,st);
   }

   /*
    * Since there are many processor cores (CPUs) in a node, getNextJob() 
    *  examines all CPU queues of a node to find a suitable job.
    * Assuming that the selected job will successfully be dispatched so that
    *  we can saftly remove it from the job queue.
    */
   public synchronized Job getNextJob(String node) {
      Job selectedJob = null;
      Set<String> cpuIds = nodeQueueMap.get(node);
      if (cpuIds == null) return null;

      int minStartTime = Integer.MAX_VALUE;
      List<Job> minJobQueue = null;
      for (String cpuId : cpuIds) {
         List<Job> jobQueue = jobQueueMap.get(cpuId);
         if (jobQueue != null) {
            if (jobQueue.size() > 0) {
               Job job = jobQueue.get(0); 
               int st = expectedStartTimeMap.get(job);
               if (st < minStartTime) {
                  minStartTime = st;
                  minJobQueue = jobQueue;
                  selectedJob = job;
               }
            }
         }
      }
      return selectedJob;
   }

   public synchronized int getExpectedStartTime(Job job) {
      Integer t = expectedStartTimeMap.get(job);
      if (t == null) return Integer.MAX_VALUE;
      return (int) t;
   }

   public synchronized void removeJob(String node, Job job) {
      Set<String> cpuIds = nodeQueueMap.get(node);
      if (cpuIds == null) return;
      for (String cpuId : cpuIds) {
         List<Job> jobQueue = jobQueueMap.get(cpuId);
         if (jobQueue != null) {
            jobQueue.remove(job);
            expectedStartTimeMap.remove(job);
            break;  // assume one job per node
         }
      }

   }

   public synchronized void checkpoint() {
     // lightweight way. It works since we append job to the queue only.
     savedJobQueueSizeMap = new HashMap<String,Integer>();
     for (String cpuId : jobQueueMap.keySet()) {
        List<Job> jobQueue = jobQueueMap.get(cpuId);
        savedJobQueueSizeMap.put(cpuId,jobQueue.size()); 
     }
     savedExpectedStartTimeMap = new HashMap<Job,Integer>(expectedStartTimeMap); 
   }

   /*
    * Truncate the job queue of each CPU according to the last checkpoint.
    * Remove the expected start time.
    */
   public synchronized void rollback() {
     for (String cpuId : jobQueueMap.keySet()) {
        List<Job> jobQueue = jobQueueMap.get(cpuId);
        Integer savedSize = savedJobQueueSizeMap.get(cpuId); 
        if (savedSize != null) {
           jobQueue.subList(savedSize,jobQueue.size()).clear();
        } else {
           jobQueue.clear();
        }
     }
     expectedStartTimeMap = savedExpectedStartTimeMap;
   }

   public String toString() {
      DateFormat dateFormat = DateFormat.getDateTimeInstance(DateFormat.SHORT,DateFormat.SHORT);
      String str = "";
      str += "----------- Schedule ---------\n";
      for (String node : nodeQueueMap.keySet()) {
         Set<String> cpus = nodeQueueMap.get(node);
         for (String cpu : cpus) {
            List<Job> jobQueue = jobQueueMap.get(cpu);
            str += cpu + ": ";
            for (Job job : jobQueue) {
               if (job.name == null) {
                  str += "Unknown PBS JOB ID ";
               }
               //str += job.name.substring(0,job.name.indexOf('.')) + "(" + expectedStartTimeMap.get(job) + ") ";
               str += job.name.substring(0,job.name.indexOf('.')) + "(" + dateFormat.format(new Date(((long) expectedStartTimeMap.get(job))*1000)) + ") ";
            }
            str += "\n";
         }
      } 
      str += "------------------------------\n";
      return str;
   }

}
